import numpy as np
import matplotlib.pyplot as plt
import animatplot as amp

plt.style.use(['ggplot', 'fast'])

fps = 30

x = np.linspace(-2, 2, num=41)
y = np.linspace(-2, 2, num=41)
t = np.linspace(0, 2 * np.pi, num=fps)
x, y, t = np.meshgrid(x, y, t)
z = np.sin(x**2 + y**2 - t)
surf = z[:, 20, :]

fig, (ax1, ax2) = plt.subplots(1, 2)

fig.suptitle('Multiple blocks')

ax1.set(title='Cross section: $y=0$', xlabel='x', ylabel='z', aspect='equal')
ax2.set(title=r'$z=\sin(x^2+y^2-t)$', xlabel='x', ylabel='y', aspect='equal')

surf_block = amp.blocks.Line(x[0, :, :], surf, ax=ax1, t_axis=1)
z_block = amp.blocks.Pcolormesh(
    x[:, :, 0], y[:, :, 0], z, ax=ax2, t_axis=2)
plt.colorbar(z_block.quad)

timeline = amp.Timeline(t, units='s', fps=fps)

ani = amp.Animation([surf_block, z_block], timeline)
ani.controls({'text': 'TIME', 'color': 'gold', 'valfmt': '%1.0f'})

ani.save_gif('Images/Multiple blocks')
plt.show()
